import scipy.io
import numpy as np
import matplotlib.pyplot as plt
import torch


# 读取.mat文件
# mat = scipy.io.loadmat('./Data/datasets/EEG/02010002rest 20150416 1017..mat')
# mat = scipy.io.loadmat('./Data/datasets/fMRI/sim4.mat')
#
# # 打印所有的键（变量名）
# print(mat.keys())
#
# print(mat)
#
# # # 假设你的数据在变量名 'eeg_data' 中
# eeg_data = mat['a02010002rest_20150416_1017mat']
# eeg_data = mat['ts']
#
# # 打印数据形状
# print('EEG Data Shape:', eeg_data.shape)

# # 打印前几个数据点
# print('EEG Data (first 10 points):', eeg_data[:10])
#
# # 画图展示EEG数据
# plt.figure(figsize=(10, 5))
#
# # 展示第一个通道的数据
# plt.plot(eeg_data[0, :])
# plt.title('EEG Data (Channel 1)')
# plt.xlabel('Sample Index')
# plt.ylabel('Amplitude')
# plt.show()
#
# # 如果想要展示多个通道的数据
# plt.figure(figsize=(10, 5))
#
# for i in range(5):  # 展示前五个通道的数据
#     plt.plot(eeg_data[i, :] + i * 10, label=f'Channel {i+1}')  # 加上偏移量以区分通道
#
# plt.title('EEG Data (First 5 Channels)')
# plt.xlabel('Sample Index')
# plt.ylabel('Amplitude')
# plt.legend(loc='upper right')
# plt.show()


# 加载 .npy 数据
# file_path = './OUTPUT/etth/samples/etth_ground_truth_24_train.npy'  # 替换为你的 .npy 文件路径
# file_path = './OUTPUT/eeg_test/ddpm_fake_eeg_test_inverse.npy'
# data = np.load(file_path)
#
# # 查看数据的形状
# print('Data shape:', data.shape)
#
# # 假设数据形状为 (channels, time_points)，选择一些通道进行绘制
# num_channels_to_plot = 5  # 你可以根据需要调整这个值
# time_points = np.arange(data.shape[1])
#
# plt.figure(figsize=(15, 10))
#
# for i in range(num_channels_to_plot):
#     plt.subplot(num_channels_to_plot, 1, i + 1)
#     plt.plot(time_points, data[i])
#     plt.title(f'Channel {i + 1}')
#     plt.xlabel('Time Points')
#     plt.ylabel('Amplitude')
#
# plt.tight_layout()
# plt.show()




def load_data(root_dir):
    import scipy.io
    import glob
    import numpy as np

    mat_files = glob.glob(f"{root_dir}/*.mat")  # 获取所有 .mat 文件路径
    data_list = []
    label_list = []

    # 目标长度
    target_length = 83876

    for mat_file in mat_files:
        mat_data = scipy.io.loadmat(mat_file)
        eeg_key = [key for key in mat_data.keys() if key.startswith('a020')][0]
        eeg_data = mat_data[eeg_key][:-1, :]  # 去掉最后一行（参考电极）

        # 如果数据的时间维度不等于目标长度，进行填充或截断
        if eeg_data.shape[1] < target_length:
            # 填充到目标长度
            pad_width = target_length - eeg_data.shape[1]
            eeg_data = np.pad(eeg_data, ((0, 0), (0, pad_width)), mode='constant')
        elif eeg_data.shape[1] > target_length:
            # 截断到目标长度
            eeg_data = eeg_data[:, :target_length]

        # 转换字节序为本地系统支持的格式
        if not eeg_data.dtype.isnative:
            eeg_data = eeg_data.byteswap().newbyteorder()

        # 根据文件名确定标签
        label = 0 if "0201" in mat_file else 1  # 抑郁患者为 0，健康为 1
        data_list.append(eeg_data)
        label_list.append(label)

        # 检查所有样本的维度
        print("检查所有样本的维度：")
        for i, data in enumerate(data_list):
            print(f"样本 {i + 1} 维度: {data.shape}")

    # 拼接为完整的数据集
    X = torch.stack([torch.tensor(data, dtype=torch.float32) for data in data_list], dim=0)  # Shape: (N_samples, N_channels, target_length)
    Y = torch.tensor(label_list, dtype=torch.int64)  # Shape: (N_samples,)

    return X, Y



load_data("./dataset")